from script.downloader.Downloader import Downloader
from typing import Callable, Union, Tuple
import os
import re
from script.Http import Http
from script.Http import HttpContext
from time import time
import json
from script.util.RawFile import RawFile


class RosefileDownloader(Downloader):
    _URL_PREFIX = 'https://rosefile.net/'
    _PATTERN_FILE_ID = re.compile(r'[\s\S]*add_ref\((?P<file_id>\d+)\);[\s\S]*')
    _PATTERN_DOWNLOAD_URL = re.compile(r'[\s\S]*<a\s+href="(?P<url>[^"]+)"[\s\S]*')
    _URL_AJAX = 'https://rosefile.net/ajax.php'

    def __init__(self, url: str, output_dir: str, filename: Union[None, str]):
        super().__init__(RosefileDownloader._name_from_url(url),
                         url)

        self._http_context = HttpContext(header={'Referer': self.url()})
        self._http = Http()
        self._output = os.path.join(output_dir, filename) if filename else os.path.join(output_dir, self.name())

    @staticmethod
    def check_url(url: str):
        return url.startswith(RosefileDownloader._URL_PREFIX)

    @staticmethod
    def create(url: str, output: str, filename: Union[None, str] = None):
        return RosefileDownloader(url, output, filename) if RosefileDownloader.check_url(url) else None

    @staticmethod
    def _name_from_url(url: str) -> str:
        return os.path.basename(url)[0:-5]

    @staticmethod
    def _filename_from_url(url: str) -> str:
        return 'rosefile_%s' % url[len(RosefileDownloader._URL_PREFIX):-5].replace('/', '_')

    def _fetch_file_id(self) -> Union[None, str]:
        # get file_id from url
        code, content = self._http_get(self.url())
        if code != 200:
            self._last_error = content
            return None
        matched = RosefileDownloader._PATTERN_FILE_ID.fullmatch(content)
        if not matched:
            self._set_last_error('invalid content format')
            return None
        return matched['file_id']

    def _fetch_wait_seconds(self) -> int:
        code, content = self._http_get('https://rosefile.net/ajax.php?action=load_time&ctime=%d' % int(time() * 1000))
        if code != 200:
            self._last_error = content
            return -1
        obj = json.loads(content)
        if obj is None or not obj['success']:
            self._last_error = content
            return -1
        return obj['waittime_s']

    def _fetch_download_url(self, file_id: str) -> Union[None, str]:
        data = 'action=load_down_addr1&file_id=%s' % file_id
        code, content = self._http_post(RosefileDownloader._URL_AJAX, data)
        if code != 200:
            self._last_error = content
            return None
        matched = RosefileDownloader._PATTERN_DOWNLOAD_URL.fullmatch(content)
        if not matched:
            self._last_error = content
            return None
        return matched['url']

    def _fetch_download_file_size(self, download_url: str) -> Union[None, int]:
        code, file_size = self._http_size(download_url)
        if code != 200:
            return None
        return file_size

    def download(self, on_progress: Union[None, Callable[[int, int], None]] = None)\
            -> Tuple[bool, Union[None, str], Union[None, str]]:
        # fetch file id
        file_id = self._fetch_file_id()
        if file_id is None:
            self._set_last_error('> fetch file_id failed: %s' % self.last_error())
            return False, None, None

        # request file address
        download_url = self._fetch_download_url(file_id)
        if download_url is None:
            self._set_last_error('> fetch download url failed: %s' % self.last_error())
            return False, None, None

        download_size = self._fetch_download_file_size(download_url)
        if download_size is None:
            self._set_last_error('> fetch download size failed: %s' % self.last_error())
            return False, None, None

        # prepare dst file dir
        dst_dir = os.path.dirname(self._output)
        if not os.path.exists(dst_dir):
            os.makedirs(dst_dir, exist_ok=True)

        # download
        with RawFile.open_file(self._output, os.O_CREAT | os.O_WRONLY) as dst_file:
            if not dst_file.is_opened():
                self._set_last_error('> open dst file failed: %s' % self.last_error())
                return False, None, None

            def on_receive(cur: int, total: int, data: bytes):
                nonlocal dst_file
                nonlocal on_progress
                nonlocal download_size

                while len(data) > 0:
                    n_written = dst_file.write(data)
                    data = data[n_written:]

                if download_size != total:
                    download_size = total

                if on_progress:
                    on_progress(cur, download_size)

            if on_progress:
                on_progress(0, download_size)
            self._http_download(download_url, offset=0, on_receive=on_receive)

            if on_progress:
                on_progress(download_size, download_size)
            return True, self._output, download_url

    def _http_get(self, _url: str) -> Tuple[int, str]:
        return self._http.get(_url, context=self._http_context)

    def _http_post(self, _url: str, _data=None) -> Tuple[int, str]:
        return self._http.post(_url, _data, context=self._http_context)

    def _http_download(self,
                       _url: str,
                       offset: int = 0,
                       length: int = -1,
                       on_receive: Union[None, Callable[[int, int, bytes], None]] = None) -> Tuple[int, str]:
        if length < 0:
            self._http_context.set_header('Range', 'bytes=%d-' % offset)
        else:
            self._http_context.set_header('Range', 'bytes=%d-%d' % (offset, offset + length - 1))
        return self._http.download(_url, on_receive=on_receive, context=self._http_context)

    def _http_size(self, _url: str) -> Tuple[int, int]:
        return self._http.size(_url, context=self._http_context)
